import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import sklearn
import skimage.io
import skimage.feature
import skimage.transform
import os
from glob import glob
import re
# save the file in pickle
import pickle
import collections

# read the folders in a list
folder = os.listdir('data')
print("image folder names: ", folder)

# read all image directories in a list
all_path =[]
for f in folder:
    all_path += glob('./data/{}/*.jpg'.format(f))

# print("all_path: ",all_path)

# read all image and label
img_all = skimage.io.ImageCollection(all_path)

print("Number of all images: ", len(img_all))
# print("File name of all images: ", img_all.files)       # i.e. './data/ElephantHead\\0000000097.jpg'

# labeling
# extract text folder. We want to extract 'Elephant' from './data/ElephantHead\\0000000095.jpg'
def extract_label(text):
    try:
        # for folders contain Head suffix, i.e. PigeonHead
        label = re.search(r'./data/(.*?)Head\\',text).group(1)
    except:
        # for folders NOT contain Head suffix, i.e. Natural
        label = re.search(r'./data/(.*?)\\',text).group(1)
        
    return label.lower()

labels = list(map(extract_label,img_all.files))
# map() allows you to process and transform all the items in an iterable without using an explicit for loop,
# Go through all image files, and extract label one by one. Thus, the length of labels is same as image files.

# Here, we try to put all images into a np array.
def buffer(io):
    return io
img_all_arrs = np.array(list(map(buffer,img_all)))

print('img_all_arrs.shape: ', img_all_arrs.shape)   # Should be [num_img, width, height, channel]

data = dict()
data['description'] ='There are 20 classes and 2057 images are there. All the images are 80 x 80 (rgb)'
data['data'] = img_all_arrs
data['target'] = labels
data['labels'] = set(labels)        # 会自动除重，只剩下每个unique的label。

pickle.dump(data,open('./pickle_files/data_animals_head_20.pickle','wb'))

count_values = collections.Counter(data['target'])
print("count_values: ",count_values)

# Plot each label with number of images inside.
y_axis = list(count_values.keys())
values = count_values.values()
plt.figure(figsize=(10,6))
plt.barh(y_axis,values,color='#C723FF')
plt.xlabel('count')
plt.title('Label Count')
plt.show()

# Plot image for each label.
plt.figure(figsize=(12,6))
for i,c in enumerate(data['labels']):
    index = data['target'].index(c)
    img = data['data'][index]
    plt.subplot(3,7,i+1)
    plt.imshow(img)
    plt.xticks([]), plt.yticks([])
    plt.title(c)
plt.show()

